"""
增强版用户类，完全复现mysql_generator.py和fake_user_simulator.py的数据生成逻辑
"""
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import List, Dict, Optional
import random
import string

try:
    from .enhanced_models import (
        UserRecord, UserRecordV2, UserGenerationParams, TrendType, AgeGroup, 
        FluctuationLabel, GenderType, ProvinceType,
        GridInfo,
        UserInfo, UserInfoWithCalulate,
        DayUserGenerator,
    )
except ImportError:
    from enhanced_models import (
        UserRecord, UserRecordV2, UserGenerationParams, TrendType, AgeGroup, 
        FluctuationLabel, GenderType, ProvinceType,
        GridInfo,
        UserInfo,
        DayUserGenerator,
    )

random.seed(42)

class EnhancedUser:
    """增强版用户类，生成符合user_ck_schema.csv规范的完整数据"""
    
    def __init__(self, params: UserGenerationParams):
        """
        初始化用户
        
        Args:
            params: 用户生成参数
        UserGenerationParams 是一个用于描述单个用户生成模拟数据所需参数的类（通常为 Pydantic 的 BaseModel 子类）。
        主要字段包括：
        - basic_info: 用户基础信息（如USER_ID、PERSON_ID、性别、年龄段、城市、区县、趋势类型等）
        - age_base_params: 年龄段基础参数（如套餐流量、通话时长、增长率、下降率、套餐金额等）
        - num_months: 需要生成的月份数
        - start_date: 起始日期（datetime类型，表示第一个月的时间）
        - individual_factor: 个体差异因子（用于调整用户行为的个体化波动）
        - grid_info: 网格ID到经纬度的映射（模拟地理位置）
        - base_info: 基站ID到经纬度的映射（模拟基站位置）
        这些参数共同决定了用户的模拟数据生成逻辑和地理属性。
        """
        self.params = params
        self.user_records: List[UserRecord] = []
        self.is_diff = random.choices([0, 1], weights=[70, 30])[0]
        
    
    def generate_user_records(self) -> List[UserRecord]:
        """
        根据用户参数，生成一个用户的所有月份的完整记录数据
        
        Returns:
            List[UserRecord]: 用户记录列表
        """
        # 生成月份列表
        months = self._generate_months()

        
        # 获取基础参数
        basic_info = self.params.basic_info
        age_params = self.params.age_base_params

        
        # 根据趋势类型设置参数
        trend_params = self._get_trend_params(basic_info.trend_type, age_params)
        
        user_records = []
        
        for month_idx, month in enumerate(months):
            record = self._generate_single_record(
                month_idx, month, basic_info, age_params, trend_params
            )
            user_records.append(record)
        
        
        self.user_records = user_records
        return user_records
    
    def _generate_months(self) -> List[str]:
        """生成月份列表"""
        end_date = self.params.start_date
        days = [
            (end_date - timedelta(days=i)).strftime("%Y%m%d") 
            for i in range(self.params.num_months-1, -1, -1)
        ]
        print(f"end_date: {end_date}, days: {days}")
        return days
    
    def _get_trend_params(self, trend_type: TrendType, age_params) -> Dict:
        """根据趋势类型获取参数调整"""
        if trend_type == TrendType.GROWTH:
            return {
                'traffic_growth_rate': age_params.growth * 1.5,
                'call_decline_rate': age_params.decline * 0.5,
                'fluctuation_level': 0.04,
                'package_adjust_rate': 0.12,
                'package_adjust_period': 2,
                'additional_service_range': (10, 15),
                'value_added_prob': 0.5
            }
        elif trend_type == TrendType.STABLE:
            return {
                'traffic_growth_rate': age_params.growth * 0.2,
                'call_decline_rate': age_params.decline * 0.2,
                'fluctuation_level': 0.02,
                'package_adjust_rate': 0.05,
                'package_adjust_period': 12,
                'additional_service_range': (8, 12),
                'value_added_prob': 0.1
            }
        elif trend_type == TrendType.DECLINE:
            return {
                'traffic_growth_rate': -age_params.growth * 1.5,
                'call_decline_rate': age_params.decline * 3,
                'fluctuation_level': 0.05,
                'package_adjust_rate': -0.08,
                'package_adjust_period': 3,
                'additional_service_range': (3, 8),
                'value_added_prob': 0.02
            }
        else:  # FLUCTUATION
            return {
                'traffic_growth_rate': age_params.growth * np.random.choice([0.5, 1.5]),
                'call_decline_rate': age_params.decline * np.random.choice([0.5, 1.5]),
                'fluctuation_level': 0.06,
                'package_adjust_rate': 0.1,
                'package_adjust_period': 3,
                'additional_service_range': (5, 15),
                'value_added_prob': 0.3
            }
    
    def _generate_single_record(
        self, month_idx: int, month: str, basic_info, age_params, trend_params
    ) -> UserRecord:
        """生成单条记录"""
        offset = month_idx
        
        # 生成是否异网用户标签（复制mysql_generator.py逻辑）
        is_diff = self.is_diff
        
        # 根据年龄段转换为具体年龄
        age = self._age_group_to_age(basic_info.age_group)
        
        # 网格和基站信息（只有非异网用户才有）

        grid_id, grid_lat, grid_long = self.params.grid_mapping.GRID_ID, self.params.grid_mapping.GRID_LAT, self.params.grid_mapping.GRID_LONG
        
        day_base_id = self.params.base_mapping.GRID_ID
        day_base_lat = self.params.base_mapping.GRID_LAT
        day_base_long = self.params.base_mapping.GRID_LONG
        
        night_base_id = self.params.base_mapping.GRID_ID
        night_base_lat = self.params.base_mapping.GRID_LAT
        night_base_long = self.params.base_mapping.GRID_LONG
        
        # 计算各项费用（只有非异网用户才计算）
        if is_diff == 0:
            package_fee = self._calculate_package_fee(offset, age_params.package, trend_params)
            traffic_usage_gb = self._calculate_traffic_usage(offset, age_params.traffic, trend_params, basic_info.trend_type)
            call_duration_minutes = self._calculate_call_duration(offset, age_params.call, trend_params)
            
            # 各项费用计算
            local_call_fee, long_distance_call_fee = self._calculate_call_fees(call_duration_minutes)
            domestic_data_fee, international_data_fee = self._calculate_data_fees(traffic_usage_gb, basic_info.age_group)
            additional_service_fee = self._calculate_additional_service_fee(trend_params)
            value_added_service_fee = self._calculate_value_added_service_fee(trend_params)
            
            # 总费用
            monthly_consumption = (
                package_fee + local_call_fee + long_distance_call_fee + 
                domestic_data_fee + international_data_fee + 
                additional_service_fee + value_added_service_fee
            )
            
            # ARPU和流量（KB）
            arpu = monthly_consumption
            flux_used = traffic_usage_gb * 1024 * 1024  # 转换为KB
            call_duration_seconds = call_duration_minutes * 60  # 转换为秒
            
        else:
            # 异网用户所有指标置空
            package_fee = None
            traffic_usage_gb = None
            call_duration_minutes = None
            local_call_fee = None
            long_distance_call_fee = None
            domestic_data_fee = None
            international_data_fee = None
            additional_service_fee = None
            value_added_service_fee = None
            monthly_consumption = None
            arpu = None
            flux_used = None
            call_duration_seconds = None
        
        # 生成社交圈ID
        family_circle_id = self._generate_random_string("FC", 3)
        social_circle_id = self._generate_random_string("SC", 3)
        community_circle_id = self._generate_random_string("CYC", 3)
        campus_circle_id = self._generate_random_string("CPC", 3)
        commute_circle_id = self._generate_random_string("CTC", 3)
        business_circle_id = self._generate_random_string("BC", 3)
        online_circle_id = self._generate_random_string("OC", 3)
        
        # 生成客户编码
        cust_id = self._generate_random_string("C", 5) if is_diff == 0 else None
        
        # 生成标签（只有非异网用户才有）
        label_is_gsm = random.randint(0, 1) if is_diff == 0 else None
        label_is_mid = random.randint(0, 1) if is_diff == 0 else None
        label_is_prem = random.randint(0, 1) if is_diff == 0 else None
        label_is_elder = random.randint(0, 1) if is_diff == 0 else None
        label_is_int = random.randint(0, 1) if is_diff == 0 else None
        
        return UserRecord(
            STATIS_TIME=month,  # 添加日期
            PROV=ProvinceType.BEIJING,
            PERSON_ID=basic_info.PERSON_ID,
            PHONE_NUM=self._generate_phone_number(),
            USER_ID=basic_info.USER_ID if is_diff == 0 else None,
            CUST_ID=cust_id,
            
            # 属性字段
            ATTR_AGE=age if is_diff == 0 else None,
            ATTR_GENDER=basic_info.ATTR_GENDER if is_diff == 0 else None,
            ATTR_PROVINCE="北京市",
            ATTR_CITY=basic_info.district,  # 使用区作为ATTR_CITY
            ATTR_COUNTY=basic_info.ATTR_COUNTY,
            
            # 网格和基站
            ATTR_GRID_ID=grid_id,
            ATTR_GRID_LAT=grid_lat,
            ATTR_GRID_LONG=grid_long,
            ATTR_DAY_BASE_ID=day_base_id,
            ATTR_DAY_BASE_LAT=day_base_lat,
            ATTR_DAY_BASE_LONG=day_base_long,
            ATTR_NIGHT_BASE_ID=night_base_id,
            ATTR_NIGHT_BASE_LAT=night_base_lat,
            ATTR_NIGHT_BASE_LONG=night_base_long,

            # 业务属性
            ATTR_FEE_FLUCTUATION=basic_info.trend_type,
            # 标签字段
            LABEL_IS_GSM=label_is_gsm,
            LABEL_IS_MID=label_is_mid,
            LABEL_IS_PREM=label_is_prem,
            LABEL_IS_DIFF=is_diff,
            LABEL_IS_ELDER=label_is_elder,
            LABEL_IS_INT=label_is_int,
        
            # 社交圈信息
            ATTR_FAMILY_CIRCLE_ID=family_circle_id,
            ATTR_SOCIAL_CIRCLE_ID=social_circle_id,
            ATTR_COMMUNITY_CIRCLE_ID=community_circle_id,
            ATTR_CAMPUS_CIRCLE_ID=campus_circle_id,
            ATTR_COMMUTE_CIRCLE_ID=commute_circle_id,
            ATTR_BUSINESS_CIRCLE_ID=business_circle_id,
            ATTR_ONLINE_CIRCLE_ID=online_circle_id,
            
            # 指标
            METRIC_ARPU=arpu,
            METRIC_FLUX_USED=flux_used,

            # 家宽属性
            METRIC_BROADBAND_COUNT=random.randint(0, 3),
            METRIC_BROADBAND_TRAFFIC=random.uniform(100000, 500000),
            METRIC_BROADBAND_ONILNETIME=random.uniform(10000, 86400),

            METRIC_MONTHLY_CONSUMPTION=monthly_consumption,
            METRIC_BASIC_PACKAGE_FEE=package_fee,
            METRIC_LOCAL_CALL_FEE=local_call_fee,
            METRIC_LONG_DISTANCE_CALL_FEE=long_distance_call_fee,
            METRIC_DOMESTIC_DATA_FEE=domestic_data_fee,
            METRIC_INTERNATIONAL_DATA_FEE=international_data_fee,
            METRIC_ADDITIONAL_SERVICE_FEE=additional_service_fee,
            METRIC_VALUE_ADDED_SERVICE_FEE=value_added_service_fee,
            METRIC_DATA_USAGE_FEE=traffic_usage_gb,
            METRIC_CALL_DURATION_MINUTES=call_duration_minutes
            
        )
    
    def _generate_random_string(self, prefix: str, length: int, start: int = 1) -> str:
        """生成随机字符串"""
        return f"{prefix}{str(start + random.randint(0, int('9' * length))).zfill(length)}"
    
    def _generate_phone_number(self) -> str:
        """生成手机号码"""
        prefixes = ["139", "138", "137", "136", "135", "134", "159", "158", "157", "150"]
        prefix = random.choice(prefixes)
        suffix = ''.join([str(random.randint(0, 9)) for _ in range(8)])
        return prefix + suffix
    
    def _age_group_to_age(self, age_group: AgeGroup) -> int:
        """将年龄段转换为具体年龄"""
        age_ranges = {
            AgeGroup.UNDER_18: (16, 17),
            AgeGroup.AGE_18_25: (18, 25),
            AgeGroup.AGE_26_30: (26, 30),
            AgeGroup.AGE_31_35: (31, 35),
            AgeGroup.AGE_35_45: (35, 45),
            AgeGroup.AGE_46_60: (46, 60),
            AgeGroup.OVER_60: (60, 80)
        }
        min_age, max_age = age_ranges[age_group]
        return random.randint(min_age, max_age)
    
    def _calculate_package_fee(self, offset: int, base_package: float, trend_params: Dict) -> float:
        """计算套餐费用"""
        if trend_params['package_adjust_rate'] > 0:
            package_adjust = 1 + trend_params['package_adjust_rate'] * (offset // trend_params['package_adjust_period'])
        else:
            package_adjust = 1 + trend_params['package_adjust_rate'] * (offset // trend_params['package_adjust_period'])
            
        if trend_params.get('package_adjust_rate', 0) < 0 and 'FLUCTUATION' in str(type(self.params.basic_info.trend_type)):
            package_adjust *= np.random.choice([-1, 1], p=[0.3, 0.7])
            
        package_fee = base_package * package_adjust * self.params.individual_factor
        return max(package_fee, base_package * 0.3)  # 最低不低于原价30%
    
    def _calculate_traffic_usage(
        self, offset: int, base_traffic: float, trend_params: Dict, trend_type: TrendType
    ) -> float:
        """计算流量使用量"""
        # 趋势增长
        traffic_growth = (1 + trend_params['traffic_growth_rate']) ** offset
        
        # 下降型用户的随机减少事件
        traffic_reduction = 1.0
        if trend_type == TrendType.DECLINE and np.random.random() < 0.3:
            traffic_reduction = np.random.uniform(0.7, 0.8)
        
        # 计算流量
        traffic = (base_traffic * traffic_growth * 
                  self.params.individual_factor * traffic_reduction)
        
        # 添加波动
        traffic *= (1 + np.random.uniform(-trend_params['fluctuation_level'], 
                                        trend_params['fluctuation_level']))
        
        return max(traffic, 0)
    
    def _calculate_call_duration(
        self, offset: int, base_call: float, trend_params: Dict
    ) -> int:
        """计算通话时长（分钟）"""
        # 通话减少趋势
        call_decline = (1 - trend_params['call_decline_rate']) ** offset
        
        # 计算通话时长
        call = (base_call * call_decline * self.params.individual_factor)
        
        # 添加波动
        call *= (1 + np.random.uniform(-trend_params['fluctuation_level'], 
                                     trend_params['fluctuation_level']))
        
        return max(int(call), 0)
    
    def _calculate_call_fees(self, call_duration: int) -> tuple:
        """计算通话费用"""
        local_call_ratio = np.random.uniform(0.6, 0.8)
        local_call_fee = call_duration * local_call_ratio * 0.1
        long_distance_call_fee = call_duration * (1 - local_call_ratio) * 0.3
        return local_call_fee, long_distance_call_fee
    
    def _calculate_data_fees(self, traffic_usage: float, age_group: AgeGroup) -> tuple:
        """计算流量费用"""
        domestic_data = traffic_usage * np.random.uniform(0.95, 1.0)
        
        # 国内流量费
        if age_group not in [AgeGroup.UNDER_18, AgeGroup.OVER_60]:
            domestic_data_fee = max(0, (domestic_data - 100) * 1)
        else:
            domestic_data_fee = 0
        
        # 国际流量费
        international_data_fee = 0
        if (np.random.random() < 0.05 and 
            age_group not in [AgeGroup.UNDER_18, AgeGroup.OVER_60]):
            international_data_fee = (traffic_usage - domestic_data) * 10
        
        return domestic_data_fee, international_data_fee
    
    def _calculate_additional_service_fee(self, trend_params: Dict) -> float:
        """计算附加业务费"""
        min_fee, max_fee = trend_params['additional_service_range']
        return np.random.uniform(min_fee, max_fee)
    
    def _calculate_value_added_service_fee(self, trend_params: Dict) -> float:
        """计算增值业务费"""
        if np.random.random() < trend_params['value_added_prob']:
            return np.random.uniform(10, 50)
        return 0
    
    
    def to_dataframe(self) -> pd.DataFrame:
        """将用户记录转换为DataFrame"""
        if not self.user_records:
            self.generate_user_records()
        
        # 转换为字典列表
        data_dicts = []
        for record in self.user_records:
            data_dict = record.dict()
            # 处理枚举类型
            for key, value in data_dict.items():
                if hasattr(value, 'value'):
                    data_dict[key] = value.value
            data_dicts.append(data_dict)
        
        return pd.DataFrame(data_dicts)



class EnhancedUserNew:
    """增强版用户类，生成符合user_ck_schema.csv规范的完整数据"""
    
    def __init__(self, params: UserInfoWithCalulate, end_date: datetime, num_days: int):
        """
        初始化用户
        
        Args:
            params: 用户生成参数。
        """
        self.params = params
        self.user_records: List[UserRecord] = []
        self.end_date = end_date
        self.num_days = num_days

    def _generate_days(self) -> List[str]:
        """生成月份列表"""
        end_date = self.end_date
        num_days = self.num_days
        days = [
            (end_date - timedelta(days=i)) 
            for i in range(num_days-1, -1, -1)
        ]
        print(f"end_date: {end_date}, days: {days}")
        return days

    def _get_trend_params(self, trend_type: TrendType, age_params = None) -> Dict:
        """根据趋势类型获取参数调整"""
        # 如果age_params为None，使用0.8到1.5的随机数
        if age_params is None:
            random_factor = np.random.uniform(0.8, 1.5)
            age_params = type('MockAgeParams', (), {
                'growth': random_factor,
                'decline': random_factor * 0.1  # decline通常是growth的10%
            })()
        
        if trend_type == TrendType.GROWTH:
            return {
                'traffic_growth_rate': age_params.growth * 1.5,
                'call_decline_rate': age_params.decline * 0.5,
                'fluctuation_level': 0.04,
                'package_adjust_rate': 0.12,
                'package_adjust_period': 2,
                'additional_service_range': (10, 15),
                'value_added_prob': 0.5
            }
        elif trend_type == TrendType.STABLE:
            return {
                'traffic_growth_rate': age_params.growth * 0.2,
                'call_decline_rate': age_params.decline * 0.2,
                'fluctuation_level': 0.02,
                'package_adjust_rate': 0.05,
                'package_adjust_period': 12,
                'additional_service_range': (8, 12),
                'value_added_prob': 0.1
            }
        elif trend_type == TrendType.DECLINE:
            return {
                'traffic_growth_rate': -age_params.growth * 1.5,
                'call_decline_rate': age_params.decline * 3,
                'fluctuation_level': 0.05,
                'package_adjust_rate': -0.08,
                'package_adjust_period': 3,
                'additional_service_range': (3, 8),
                'value_added_prob': 0.02
            }
        else:  # FLUCTUATION
            return {
                'traffic_growth_rate': age_params.growth * np.random.choice([0.5, 1.5]),
                'call_decline_rate': age_params.decline * np.random.choice([0.5, 1.5]),
                'fluctuation_level': 0.06,
                'package_adjust_rate': 0.1,
                'package_adjust_period': 3,
                'additional_service_range': (5, 15),
                'value_added_prob': 0.3
            }

    def generate_user_records(self) -> List[UserRecord]:
        """
        根据用户参数，生成一个用户的所有月份的完整记录数据
        
        Returns:
            List[UserRecord]: 用户记录列表
        """
        # 生成月份列表
        days = self._generate_days()

        
        # 根据趋势类型设置参数
        trend_params = self._get_trend_params(self.params.trend_type)
        
        user_records = []
        
        day_gen = DayUserGenerator(days, self.params)
        for record in day_gen.generator():
            user_records.append(record)
        
        
        self.user_records = user_records
        return user_records